import cv2
import numpy as np


def show(img, name=""):
    cv2.imshow(name, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


paper = cv2.imread("rsc/paper.jpg")
paper_copy = paper.copy()
gray = cv2.cvtColor(paper, cv2.COLOR_BGR2GRAY)

# 找到大框
thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1]
show(thresh, "thresh")
top_hat = cv2.morphologyEx(thresh, cv2.MORPH_TOPHAT, np.ones((9, 3)), iterations=1)
show(top_hat, "top_hat")

contours = cv2.findContours(top_hat, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
contours = sorted(contours, key=cv2.contourArea, reverse=True)[:5]
cv2.drawContours(paper_copy, contours, -1, (0, 255, 0), 2)
show(paper_copy, "contours")

# 找到选择题框
targets = []
for contour in contours:
    length = cv2.arcLength(contour, closed=True)
    approx = cv2.approxPolyDP(contour, 0.05 * length, closed=True)
    if len(approx) == 4:
        rect = cv2.boundingRect(contour)
        x, y, w, h = rect
        if w * 0.8 > h > w * 0.5:
            targets.append(rect)

for (x, y, w, h) in targets:
    cv2.rectangle(paper_copy, (x, y), (x + w, y + h), (255, 0, 0), 2)
show(paper_copy, "choices")


def sort_rect(group_rects, r=5):
    """
    将表格状分布的矩形框按从左到右，从上到下排序，当两个矩形框的y偏差在r容忍范围内时视为矩形框在同一行
    :param group_rects: 待排序的矩形框列表
    :param r: 属于同一行的容忍范围
    :return: 排序好的矩形框列表
    """
    group_rects = sorted(group_rects, key=lambda g: g[1])
    groups = []
    group = [group_rects[0]]
    for i in range(1, len(group_rects)):
        if -r / 2 < (group_rects[i - 1][1] - group_rects[i][1]) < r / 2:
            # y差距小说明是同一行
            group.append(group_rects[i])
        else:
            # 不在同一行了，将存下的同一行的方框按x排序好
            groups += sorted(group, key=lambda g: g[0])
            # 从新的一行开始
            group = [group_rects[i]]
    groups += sorted(group, key=lambda g: g[0])  # 最后一行不会进else
    return groups


# 和原图不一样选项的自定义答案，用于测试
correct_answers = [
    0, 0, 1, 1, 2,
    0, 1, 1, 2, 0,
    0, 1, 2, 0, 2,  # 0, 2, 2, 1, 2,  改了2个
    1, 2, 1, 0, 0,
    1, 2, 3, 0, 2,
    0, 3, 1, 1, 2,  # 0, 3, 0, 1, 3,  改了2个
    1, 2, 1, 3, 0,
    3, 1, 3, 0, 4,  # 6, 1, 3, 0, 4,  改了1个
    1, 0, 1, 3, 0,
    1, 3, 2, 3, 3,
    1, 1, 3, 2, 2
]
answers = []
task_count = 0
correct_count = 0
final_img = paper.copy()

# 只有一个选择题框
if len(targets) == 1:
    X, Y, W, H = targets[0]
    # 使用mask而不是直接裁剪是为了方便在原图找到要画方框的位置
    mask = np.zeros_like(gray, dtype=np.int8)
    mask[Y + 2:Y + H - 2, X + 2:X + W - 2] = 255  # 小一点去掉边界的线，后面好分题框
    choose_place = cv2.bitwise_and(paper, paper, mask=mask)
    choose_copy = choose_place.copy()
    show(choose_place, "choose_place")

    # 分割各个题框（每个5题）
    choose_gray = cv2.cvtColor(choose_place, cv2.COLOR_BGR2GRAY)
    choose_thresh = cv2.threshold(choose_gray, 230, 255, cv2.THRESH_BINARY_INV)[1]
    choose_thresh[mask == 0] = 0  # 将掩膜外的地方改回黑色
    show(choose_thresh, "choose_thresh")
    choose_dilate = cv2.dilate(choose_thresh, np.ones((3, 3)), iterations=2)
    show(choose_dilate, "choose_dilate")

    group_contours = cv2.findContours(choose_dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
    group_contours = sorted(group_contours, key=cv2.contourArea, reverse=True)[:11]
    # print(len(group_contours))
    cv2.drawContours(choose_copy, group_contours, -1, (0, 255, 0), 2)
    show(choose_copy, "group_contours")

    # 题框按题序排好序
    group_rects = [cv2.boundingRect(group_contour) for group_contour in group_contours]
    groups = sort_rect(group_rects)

    for (x, y, w, h) in groups:
        # 分割每个题目
        mask = np.zeros_like(gray, dtype=np.int8)
        mask[y:y + h, x + 7:x + w - 1] = 255  # 范围是为了去掉题号
        group_place = cv2.bitwise_and(paper, paper, mask=mask)
        show(group_place, "group_place")
        # print(group_place.shape)

        group_gray = cv2.cvtColor(group_place, cv2.COLOR_BGR2GRAY)
        group_thresh = cv2.threshold(group_gray, 100, 255, cv2.THRESH_BINARY_INV)[1]
        group_thresh[mask == 0] = 0
        # show(group_thresh)
        # 题组
        choice_shape = (8, 10)  # 一个选项的大小
        for i in range(h // choice_shape[0]):  # h // choice_shape[0] 得出题目数量
            # 每一题
            mask = np.zeros_like(gray, dtype=np.int8)
            mask[y + i * choice_shape[0]:y + (i + 1) * choice_shape[0], x + 7:x + w - 1] = 255
            task_place = cv2.bitwise_and(paper, paper, mask=mask)
            # show(task_place, "task_place")

            task_gray = cv2.cvtColor(task_place, cv2.COLOR_BGR2GRAY)
            task_thresh = cv2.threshold(task_gray, 100, 255, cv2.THRESH_BINARY_INV)[1]
            task_thresh[mask == 0] = 0
            # show(task_thresh, "task_thresh")

            choose_cont = cv2.findContours(task_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
            if not len(choose_cont) == 0 and cv2.contourArea(choose_cont[0]) > 10:
                # 有选择
                xx, yy, ww, hh = cv2.boundingRect(choose_cont[0])
                choice = (xx - x - 7) // choice_shape[1]  # xx - (x + 7)
                answers.append(choice)
                if choice == correct_answers[task_count]:
                    # 答对了，用绿色画框
                    cv2.rectangle(final_img, (xx, yy), (xx + ww, yy + hh), (0, 255, 0), 1)
                    correct_count += 1
                else:
                    # 答错了，用红色画正确位置
                    cv2.rectangle(final_img,
                                  (
                                      x + 7 + correct_answers[task_count] * choice_shape[1] + 2,  # + 2 是因为方框有点偏移
                                      y + i * choice_shape[0] + 2
                                  ),
                                  (
                                      x + 7 + correct_answers[task_count] * choice_shape[1] + 2 + ww,
                                      y + i * choice_shape[0] + 2 + hh
                                  ),
                                  (0, 0, 255), 1)
            else:
                answers.append(-1)
            task_count += 1  # 进入下一题
    print(answers, len(answers))
    # 显示正确率
    cv2.putText(final_img, f"{round(correct_count / task_count * 100, 2)}%", (X, Y), cv2.FONT_HERSHEY_SIMPLEX, 0.65,
                (0, 0, 255), 2)
    show(final_img)
    cv2.imwrite("rsc/result.jpg", final_img)
"""
print(answers)

[0, 0, 1, 1, 2,
 0, 1, 1, 2, 0, 
 0, 2, 2, 1, 2, 
 1, 2, 1, 0, 0, 
 1, 2, 3, 0, 2, 
 0, 3, 0, 1, 3, 
 1, 2, 1, 3, 0, 
 6, 1, 3, 0, 4, 
 1, 0, 1, 3, 0, 
 1, 3, 2, 3, 3, 
 1, 1, 3, 2, 2]
"""
